import pandas as pd
import numpy as np
import os
import joblib
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, roc_auc_score, confusion_matrix
from sklearn.preprocessing import StandardScaler
import sys
import matplotlib.pyplot as plt
from tqdm import tqdm

# 添加当前目录到系统路径，以便导入model_factory和B1FeatureExtractor
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from model_factory import ModelFactory
from B1_train import B1FeatureExtractor

def load_stock_data(file_path):
    """
    加载单个股票数据
    
    Parameters:
    file_path: str, 数据文件路径
    
    Returns:
    DataFrame: 股票数据
    """
    try:
        df = pd.read_csv(file_path)
        # 特征列验证
        required_columns = ['j', 'Vol_x_Price_Change', 'Relative_Volume', 'ATR',
                           'ATR_percentage', 'MA_Trend_Signal', 'Close_Change_Rate', 'Volume_Change_Rate', 'dif']
        missing_cols = [col for col in required_columns if col not in df.columns]
        if missing_cols:
            print(f"警告：文件 {file_path} 缺失必要列: {missing_cols}")
            return None
            
        df['date'] = pd.to_datetime(df['date'])
        df = df.sort_values('date').reset_index(drop=True)

        #过滤数据，只保留2021年12月1日 - 2024年12月30日的数据
        start_date = pd.to_datetime('2021-12-01')
        end_date = pd.to_datetime('2024-12-30')
        df = df[(df['date'] >= start_date) & (df['date'] <= end_date)].copy()

        return df
    except Exception as e:
        print(f"加载文件 {file_path} 时出错: {e}")
        return None

def train_individual_model(stock_file, data_dir, model_dir, sequence_length=15):
    """
    为单个股票训练模型
    
    Parameters:
    stock_file: str, 股票数据文件名
    data_dir: str, 数据目录路径
    model_dir: str, 模型保存目录路径
    sequence_length: int, 序列长度
    """
    file_path = os.path.join(data_dir, stock_file)
    
    # 加载股票数据
    print(f"开始处理股票: {stock_file}")
    df = load_stock_data(file_path)
    if df is None:
        return
    
    # 准备特征数据
    feature_extractor = B1FeatureExtractor(sequence_length=sequence_length)
    try:
        X, y = feature_extractor.prepare_features(df, target_col='B1_tag')
        if len(X) == 0 or len(y) == 0:
            print(f"股票 {stock_file} 没有足够的数据用于训练")
            return
        
        # 只保留J值小于10的样本，将标签转换为二分类
        y_binary = (y == 1).astype(int)
        
        # 特征精度控制
        X = np.array(X).round(4)
        
        # 缺失值处理增强
        if np.isnan(X).any():
            print(f"警告：股票 {stock_file} 发现特征数据包含NaN值，使用前向填充和后向填充处理")
            X = pd.DataFrame(X).ffill().bfill().fillna(0).values
        
        # 显示数据分布
        print(f"股票 {stock_file} 数据分布: B1_tag=1 的样本数: {np.sum(y_binary == 1)}, B1_tag!=1 的样本数: {np.sum(y_binary != 1)}")
        
        # 分割训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y_binary, test_size=0.2, random_state=42)
        
        # 如果训练样本太少，跳过训练
        if len(X_train) < 10:
            print(f"股票 {stock_file} 训练样本太少，跳过训练")
            return
        
        # 创建并训练模型
        print(f"开始训练股票 {stock_file} 的随机森林模型...")
        model = ModelFactory.create_model('rf', n_estimators=100, random_state=42)
        model.train(X_train, y_train)
        
        # 在测试集上进行预测
        y_pred_proba = model.predict(X_test)
        y_pred = (y_pred_proba > 0.5).astype(int)
        
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        try:
            auc_score = roc_auc_score(y_test, y_pred_proba)
        except:
            auc_score = 0.0
        
        print(f"股票 {stock_file} 模型评估结果:")
        print(f"  准确率: {accuracy:.4f}")
        print(f"  AUC Score: {auc_score:.4f}")
        
        # 保存模型和scaler
        stock_code = stock_file.split('_')[0]  # 从文件名提取股票代码
        model_path = os.path.join(model_dir, f"{stock_code}_b1_rf_model.pkl")
        scaler_path = os.path.join(model_dir, f"{stock_code}_b1_rf_scaler.pkl")
        
        print(f"保存股票 {stock_file} 的模型和scaler...")
        joblib.dump(model, model_path)
        joblib.dump(feature_extractor.scaler, scaler_path)
        
        print(f"股票 {stock_file} 的模型已保存到: {model_path}")
        print(f"股票 {stock_file} 的scaler已保存到: {scaler_path}")
        
    except Exception as e:
        print(f"处理股票 {stock_file} 时出错: {e}")

def train_all_individual_models():
    """
    为所有股票分别训练模型
    """
    # 数据目录和模型保存目录
    data_dir = os.path.join("..", "tdx_data_process")
    # 修改为与Python文件同级的目录
    model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "b1_model_every")
    
    # 检查数据目录是否存在
    if not os.path.exists(data_dir):
        print(f"数据目录不存在: {data_dir}")
        return
    
    # 创建模型保存目录
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
        print(f"创建模型保存目录: {model_dir}")
    
    # 获取所有股票数据文件
    stock_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
    
    print(f"找到 {len(stock_files)} 只股票的数据文件")
    
    # 为每只股票训练模型
    for i, stock_file in enumerate(tqdm(stock_files, desc="训练进度")):
        try:
            train_individual_model(stock_file, data_dir, model_dir)
            
            if (i + 1) % 100 == 0:
                print(f"已处理 {i + 1} 只股票")
        except Exception as e:
            print(f"处理股票文件 {stock_file} 时出错: {e}")
    
    print("所有股票的模型训练完成")

def main():
    """
    主函数
    """
    print("为每只股票单独训练B1模型")
    print("=" * 50)
    train_all_individual_models()

if __name__ == "__main__":
    main()